#include <cv_bridge/cv_bridge.h>
#include <image_transport/image_transport.h>
#include <ros/ros.h>
#include <std_msgs/Int8.h>

#include <chrono>
#include <mutex>
#include <opencv2/opencv.hpp>
#include <string>
#include <thread>

#include "flashai/commons/string_utils.h"
#include "flashros/commons/utils.h"

class RosCameraHelper {
 public:
  void Init(int argc, char** argv,
            std::string default_node_name = "two_stream_publisher",
            std::string data_root_config_name = "data_root",
            std::string config_param_base_name = "/config_filename",
            std::string input_config_section = "input",
            std::string output_config_section = "output",
            std::string channel_config_section = "channel") {
    // ROS 初始化
    // 注意，这里的 default_node_name 只有在 launch 文件中没有指定 node name
    // 的情况下才生效
    flashros::commons::RosInit(argc, argv, default_node_name);
    nh_ = new ros::NodeHandle;
    ROS_WARN("Successfully init ros");

    // 数据根目录
    nh_->param<std::string>(data_root_config_name, data_root_, "");

    // 读取配置文件
    // 配置文件默认保存在参数 `[node_name]/config_filename` 中
    std::string cur_node_name = ros::this_node::getName();
    std::string config_param_name = cur_node_name + config_param_base_name;
    cfg_reader = flashros::commons::InitConfig(*nh_, config_param_name);
    ROS_WARN_STREAM("Successfully read param " << config_param_name);

    // 相关基本配置信息
    retry_ = cfg_reader->GetBoolean(input_config_section, "retry", false);
    video_path_ = cfg_reader->GetString(input_config_section, "video_path", "");
    video_height_ =
        cfg_reader->GetInteger(input_config_section, "video_height", 0);
    video_width_ =
        cfg_reader->GetInteger(input_config_section, "video_width", 0);
    video_fps_ = cfg_reader->GetInteger(input_config_section, "video_fps", 0);

    // 构建 Publisher 相关参数
    max_fps_ = cfg_reader->GetInteger(output_config_section, "max_fps", 10);
    int topic_queue_size =
        cfg_reader->GetInteger(output_config_section, "topic_queue_size", 1);
    std::string topic_names_str =
        cfg_reader->GetString(output_config_section, "topic_names", "");
    std::vector<std::string> topic_names;
    flashai::utils::SplitStr(topic_names_str, topic_names);
    loop_rate_ = new ros::Rate(max_fps_);
    // 注意，这里要求 topic_names 字段的 topic 顺序与 channel change 中的
    // channel 编号对应
    image_transports_.resize(topic_names.size());
    publishers_.resize(topic_names.size());
    for (int i = 0; i < topic_names.size(); i++) {
      auto topic_name = topic_names[i];
      image_transports_[i] = new image_transport::ImageTransport(*nh_);
      publishers_[i] =
          image_transports_[i]->advertise(topic_name, topic_queue_size);
      ROS_WARN_STREAM("Successfully init publisher for topic " << topic_name);
    }
    ROS_WARN_STREAM("Successfully create " << publishers_.size()
                                           << " publisher(s)");

    // 初始化channel相关
    if (publishers_.size() > 1) {
      ROS_WARN_STREAM("Enable channel change topic");
      InitChennlChange(channel_config_section);
    } else {
      cur_channel_ = 0;
    }
  }

  void OpenSource(cv::VideoCapture& cap) {
    // 通过 video_path_ 初始化数据源，可能是摄像头、rtsp、本地视频文件

    try {
      int camera_id = std::stoi(video_path_);
      ROS_WARN_STREAM("Try to open webcam " << camera_id);
      cap.open(camera_id);
    } catch (...) {
      if (video_path_.find("rtsp://") != std::string::npos) {
        // rtsp
        ROS_WARN_STREAM("Try to open rtsp " << video_path_);
        cap.open(video_path_);
      } else {
        // normal video path
        auto video_path = data_root_ + video_path_;
        ROS_WARN_STREAM("Try to open normal video file " << video_path);
        flashros::commons::CheckFileExists(video_path);
        cap.open(video_path);
      }
    }

    // set videocap params
    //ROS_WARN_STREAM("Try to set input video to size ("
    //                << video_height_ << ", " << video_width_ << ") with fps "
    //                << video_fps_);
    //cap.set(cv::CAP_PROP_FRAME_WIDTH, video_width_);
    //cap.set(cv::CAP_PROP_FRAME_HEIGHT, video_height_);
    //cap.set(cv::CAP_PROP_FPS, video_fps_);
  }

  void Loop(cv::VideoCapture& cap) {
    bool read_flag = false;
    cv::Mat frame;
    sensor_msgs::ImagePtr msg;
    loop_rate_->reset();

    while (nh_->ok()) {
      // do read frame
      read_flag = cap.read(frame);
      if (!read_flag) break;

      // build msg
      msg = cv_bridge::CvImage(std_msgs::Header(), "bgr8", frame).toImageMsg();
      msg->header.stamp = ros::Time::now();

      // publish msg
      lock_.lock();
      // 选择当前 chennel 对应的 publisher
      image_transport::Publisher& cur_publisher = publishers_[cur_channel_];
      lock_.unlock();
      cur_publisher.publish(msg);
      ros::spinOnce();

      // limit max fps
      loop_rate_->sleep();
      auto duration = loop_rate_->cycleTime();

      // Log per second
      // msg format `Read frame with shape (224, 224) at 2021-09-10 18:29:30`
      ROS_WARN_STREAM_THROTTLE(
          max_fps_, "Read frame with shape ("
                        << frame.rows << ", " << frame.cols << ") within "
                        << int(duration.toSec() * 1000) << " ms at "
                        << flashros::commons::GetTimePointString());
    }  // while nh.ok()
  }

  void Run() {
    cv::VideoCapture cap;
    while (retry_ && nh_->ok()) {
      // Try to open source
      OpenSource(cap);

      if (cap.isOpened()) {
        Loop(cap);
        ROS_INFO_STREAM(video_path_ << " is closed...");
      } else {
        ROS_ERROR_STREAM("Cannot open " << video_path_);
      }

      if (retry_) {
        ROS_WARN_STREAM("Will try to reopen " << video_path_
                                              << "  in 5 secs...");
        std::this_thread::sleep_for(std::chrono::seconds(5));
      }

      cap.release();
    }
  }

  void Release() {
    for (auto& publisher : publishers_) {
      publisher.shutdown();
    }
    for (auto& image_transport : image_transports_) {
      delete image_transport;
    }
    delete nh_;
    delete cfg_reader;
    delete loop_rate_;

    // TODO: chane channel shift
    channel_listen_thread_running_ = false;
  }

 private:
  // base params
  ros::NodeHandle* nh_;
  INIReader* cfg_reader;
  bool retry_;

  // video param
  std::string video_path_, data_root_;
  int video_height_, video_width_, video_fps_;

  // publisher params
  std::vector<image_transport::ImageTransport*> image_transports_;
  std::vector<image_transport::Publisher> publishers_;
  ros::Rate* loop_rate_;
  int max_fps_;

  // two-stream related params
  std::thread channel_command_listen_thread_;
  std::string channel_shift_topic_;
  std::mutex lock_;
  int cur_channel_;
  bool channel_listen_thread_running_;

  friend void channelSwitchCallback(const std_msgs::Int8::ConstPtr&);

  void InitChennlChange(const std::string& channel_config_section) {
    // 初始化“切换通道”相关参数
    // 主要就是启动线程，监听 `channel_topic_name`，切换 channel
    channel_listen_thread_running_ = true;
    cur_channel_ =
        cfg_reader->GetInteger(channel_config_section, "default_channel", 0);
    channel_shift_topic_ =
        cfg_reader->GetString(channel_config_section, "channel_topic_name", "");
    ROS_WARN_STREAM("Channel change topic " << channel_shift_topic_);

    // auto callback_fn = [&](const std_msgs::Int8::ConstPtr& msg) {};

    channel_command_listen_thread_ = std::thread([&]() {
      auto channel_shift_listen_callback_fn =
          [&](const std_msgs::Int8::ConstPtr& msg) {
            lock_.lock();
            cur_channel_ = msg->data;
            ROS_WARN_STREAM("Change channel to " << cur_channel_);
            lock_.unlock();
          };

      ros::Subscriber sub = nh_->subscribe<std_msgs::Int8>(
          channel_shift_topic_, 1, channel_shift_listen_callback_fn);
      ROS_WARN_STREAM("ROS Init subscirbe for " << channel_shift_topic_);
      while (channel_listen_thread_running_) {
        ros::spinOnce();
      }
    });
  }
};

int main(int argc, char** argv) {
  RosCameraHelper helper;
  helper.Init(argc, argv);
  helper.Run();
  helper.Release();
  return 0;
}
